package com.flink.hbase.statemanagement;

import com.flink.hbase.statemanagement.config.StateBackendConfig;
import com.flink.hbase.statemanagement.managed.ValueStateExample;
import com.flink.hbase.statemanagement.managed.ListStateExample;
import com.flink.hbase.statemanagement.managed.MapStateExample;
import com.flink.hbase.statemanagement.managed.BroadcastStateExample;
import com.flink.hbase.statemanagement.managed.RichFunctionStateExample;
import com.flink.hbase.statemanagement.raw.RawStateExample;
import com.flink.hbase.statemanagement.checkpoint.CheckpointedFunctionExample;

import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.environment.CheckpointConfig;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.runtime.state.StateBackend;

/**
 * 状态管理主作业类
 * 
 * 整合所有状态管理示例：
 * 1. RocksDB状态后端配置
 * 2. 托管状态示例（ValueState、ListState、MapState、BroadcastState）
 * 3. 原生状态示例
 * 4. Rich Function状态示例
 * 5. CheckpointedFunction示例
 * 
 * @author Flink Team
 */
public class StateManagementJob {
    
    /**
     * 运行模式枚举
     */
    public enum RunMode {
        VALUE_STATE,           // ValueState示例
        LIST_STATE,            // ListState示例
        MAP_STATE,             // MapState示例
        BROADCAST_STATE,       // BroadcastState示例
        RAW_STATE,             // 原生状态示例
        RICH_FUNCTION,         // Rich Function状态示例
        CHECKPOINTED_FUNCTION, // CheckpointedFunction示例
        ALL_EXAMPLES           // 所有示例
    }
    
    /**
     * 环境配置枚举
     */
    public enum Environment {
        DEV,     // 开发环境
        TEST,    // 测试环境
        PROD     // 生产环境
    }
    
    /**
     * 主方法
     */
    public static void main(String[] args) throws Exception {
        // 解析命令行参数
        RunMode runMode = parseRunMode(args);
        Environment environment = parseEnvironment(args);
        String checkpointDir = parseCheckpointDir(args);
        
        // 创建执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        
        // 配置并行度
        int parallelism = parseParallelism(args);
        env.setParallelism(parallelism);
        
        // 配置状态后端
        configureStateBackend(env, environment, checkpointDir);
        
        // 根据运行模式执行相应示例
        executeByMode(env, runMode);
        
        // 执行作业
        String jobName = "状态管理示例 - " + runMode.name();
        env.execute(jobName);
    }
    
    /**
     * 解析运行模式
     */
    private static RunMode parseRunMode(String[] args) {
        for (int i = 0; i < args.length - 1; i++) {
            if ("--mode".equals(args[i])) {
                try {
                    return RunMode.valueOf(args[i + 1].toUpperCase());
                } catch (IllegalArgumentException e) {
                    System.err.println("无效的运行模式: " + args[i + 1]);
                    printUsage();
                    System.exit(1);
                }
            }
        }
        return RunMode.ALL_EXAMPLES; // 默认运行所有示例
    }
    
    /**
     * 解析环境配置
     */
    private static Environment parseEnvironment(String[] args) {
        for (int i = 0; i < args.length - 1; i++) {
            if ("--env".equals(args[i])) {
                try {
                    return Environment.valueOf(args[i + 1].toUpperCase());
                } catch (IllegalArgumentException e) {
                    System.err.println("无效的环境配置: " + args[i + 1]);
                    printUsage();
                    System.exit(1);
                }
            }
        }
        return Environment.DEV; // 默认开发环境
    }
    
    /**
     * 解析检查点目录
     */
    private static String parseCheckpointDir(String[] args) {
        for (int i = 0; i < args.length - 1; i++) {
            if ("--checkpoint-dir".equals(args[i])) {
                return args[i + 1];
            }
        }
        return "hdfs://namenode:9000/flink/checkpoints"; // 默认检查点目录
    }
    
    /**
     * 解析并行度
     */
    private static int parseParallelism(String[] args) {
        for (int i = 0; i < args.length - 1; i++) {
            if ("--parallelism".equals(args[i])) {
                try {
                    return Integer.parseInt(args[i + 1]);
                } catch (NumberFormatException e) {
                    System.err.println("无效的并行度: " + args[i + 1]);
                    printUsage();
                    System.exit(1);
                }
            }
        }
        return 2; // 默认并行度
    }
    
    /**
     * 配置状态后端
     */
    private static void configureStateBackend(StreamExecutionEnvironment env, 
                                            Environment environment, 
                                            String checkpointDir) {
        
        System.out.println("=== 状态后端配置 ===");
        System.out.println("环境: " + environment);
        System.out.println("检查点目录: " + checkpointDir);
        
        // 根据环境配置状态后端
        switch (environment) {
            case DEV:
                StateBackendConfig.configureDevRocksDB(env);
                break;
            case TEST:
                StateBackendConfig.configureBasicRocksDB(env, checkpointDir);
                break;
            case PROD:
                StateBackendConfig.configureProductionRocksDB(env, checkpointDir, 
                    "hdfs://namenode:9000/flink/savepoints");
                break;
        }
        
        // 打印状态后端配置信息
        StateBackendConfig.printStateBackendInfo(env);
    }
    
    /**
     * 根据运行模式执行相应示例
     */
    private static void executeByMode(StreamExecutionEnvironment env, RunMode runMode) throws Exception {
        System.out.println("=== 运行模式: " + runMode + " ===");
        
        switch (runMode) {
            case VALUE_STATE:
                ValueStateExample.runValueStateExample(env);
                break;
                
            case LIST_STATE:
                ListStateExample.runListStateExample(env);
                break;
                
            case MAP_STATE:
                MapStateExample.runMapStateExample(env);
                break;
                
            case BROADCAST_STATE:
                BroadcastStateExample.runBroadcastStateExample(env);
                break;
                
            case RAW_STATE:
                RawStateExample.runRawStateExample(env);
                break;
                
            case RICH_FUNCTION:
                RichFunctionStateExample.runRichFunctionStateExample(env);
                break;
                
            case CHECKPOINTED_FUNCTION:
                CheckpointedFunctionExample.runCheckpointedFunctionExample(env);
                break;
                
            case ALL_EXAMPLES:
                runAllExamples(env);
                break;
        }
    }
    
    /**
     * 运行所有示例
     */
    private static void runAllExamples(StreamExecutionEnvironment env) throws Exception {
        System.out.println("🚀 开始运行所有状态管理示例...");
        
        // 注意：在实际生产环境中，不建议在一个作业中运行所有示例
        // 这里仅用于演示目的
        
        // 1. ValueState示例
        System.out.println("\n--- ValueState示例 ---");
        ValueStateExample.runValueStateExample(env);
        
        // 2. ListState示例
        System.out.println("\n--- ListState示例 ---");
        ListStateExample.runListStateExample(env);
        
        // 3. MapState示例
        System.out.println("\n--- MapState示例 ---");
        MapStateExample.runMapStateExample(env);
        
        // 4. BroadcastState示例
        System.out.println("\n--- BroadcastState示例 ---");
        BroadcastStateExample.runBroadcastStateExample(env);
        
        // 5. Rich Function状态示例
        System.out.println("\n--- Rich Function状态示例 ---");
        RichFunctionStateExample.runRichFunctionStateExample(env);
        
        // 6. CheckpointedFunction示例
        System.out.println("\n--- CheckpointedFunction示例 ---");
        CheckpointedFunctionExample.runCheckpointedFunctionExample(env);
        
        // 7. 原生状态示例
        System.out.println("\n--- 原生状态示例 ---");
        RawStateExample.runRawStateExample(env);
        
        System.out.println("\n✅ 所有状态管理示例配置完成");
        printStateManagementSummary();
    }
    
    /**
     * 打印状态管理总结
     */
    private static void printStateManagementSummary() {
        System.out.println("\n=== 状态管理总结 ===");
        System.out.println("📊 托管状态 (Managed State):");
        System.out.println("  - ValueState: 存储单个值的状态");
        System.out.println("  - ListState: 存储列表的状态");
        System.out.println("  - MapState: 存储键值对的状态");
        System.out.println("  - BroadcastState: 广播状态，所有任务共享");
        
        System.out.println("\n🔧 Rich Function状态:");
        System.out.println("  - 通过RuntimeContext访问状态");
        System.out.println("  - 支持状态TTL配置");
        System.out.println("  - 自动参与Checkpoint机制");
        
        System.out.println("\n📸 CheckpointedFunction:");
        System.out.println("  - 自定义状态快照和恢复逻辑");
        System.out.println("  - 算子状态管理");
        System.out.println("  - 状态一致性保证");
        
        System.out.println("\n🗄️ 原生状态 (Raw State):");
        System.out.println("  - 用户自定义序列化");
        System.out.println("  - 灵活的状态存储");
        System.out.println("  - 需要手动管理");
        
        System.out.println("\n🚀 RocksDB状态后端:");
        System.out.println("  - 高性能状态存储");
        System.out.println("  - 支持大状态");
        System.out.println("  - 可配置优化参数");
        
        System.out.println("\n⚡ 性能优化建议:");
        System.out.println("  - 合理设置状态TTL");
        System.out.println("  - 使用适当的状态类型");
        System.out.println("  - 定期清理过期状态");
        System.out.println("  - 配置合适的Checkpoint间隔");
        
        System.out.println("========================");
    }
    
    /**
     * 打印使用说明
     */
    private static void printUsage() {
        System.out.println("\n=== 使用说明 ===");
        System.out.println("java -jar state-management-examples.jar [OPTIONS]");
        System.out.println("\nOPTIONS:");
        System.out.println("  --mode <MODE>           运行模式 (默认: ALL_EXAMPLES)");
        System.out.println("                          可选值: VALUE_STATE, LIST_STATE, MAP_STATE,");
        System.out.println("                                 BROADCAST_STATE, RAW_STATE, RICH_FUNCTION,");
        System.out.println("                                 CHECKPOINTED_FUNCTION, ALL_EXAMPLES");
        System.out.println("  --env <ENV>             环境配置 (默认: DEV)");
        System.out.println("                          可选值: DEV, TEST, PROD");
        System.out.println("  --checkpoint-dir <DIR>  检查点目录 (默认: hdfs://namenode:9000/flink/checkpoints)");
        System.out.println("  --parallelism <NUM>     并行度 (默认: 2)");
        System.out.println("\n示例:");
        System.out.println("  # 运行ValueState示例");
        System.out.println("  java -jar state-management-examples.jar --mode VALUE_STATE --env DEV");
        System.out.println("  # 运行所有示例，生产环境配置");
        System.out.println("  java -jar state-management-examples.jar --mode ALL_EXAMPLES --env PROD --parallelism 4");
        System.out.println("================");
    }
    
    /**
     * 状态管理最佳实践工具类
     */
    public static class StateManagementUtils {
        
        /**
         * 检查状态后端配置
         */
        public static void validateStateBackendConfig(StreamExecutionEnvironment env) {
            StateBackend stateBackend = env.getStateBackend();
            
            if (stateBackend == null) {
                System.err.println("⚠️ 警告: 未配置状态后端，使用默认配置");
                return;
            }
            
            CheckpointConfig checkpointConfig = env.getCheckpointConfig();
            
            // 检查Checkpoint配置
            if (checkpointConfig.getCheckpointInterval() <= 0) {
                System.err.println("⚠️ 警告: 未启用Checkpoint");
            }
            
            // 检查Checkpoint模式
            if (checkpointConfig.getCheckpointingMode() != CheckpointingMode.EXACTLY_ONCE) {
                System.err.println("⚠️ 警告: 建议使用EXACTLY_ONCE模式");
            }
            
            // 检查并发Checkpoint数量
            if (checkpointConfig.getMaxConcurrentCheckpoints() > 1) {
                System.err.println("⚠️ 警告: 并发Checkpoint可能影响性能");
            }
            
            System.out.println("✅ 状态后端配置检查完成");
        }
        
        /**
         * 获取状态管理建议
         */
        public static void printStateManagementTips() {
            System.out.println("\n=== 状态管理最佳实践 ===");
            System.out.println("1. 🎯 选择合适的状态类型:");
            System.out.println("   - 单个值 -> ValueState");
            System.out.println("   - 列表数据 -> ListState");
            System.out.println("   - 键值对 -> MapState");
            System.out.println("   - 广播配置 -> BroadcastState");
            
            System.out.println("\n2. ⏰ 配置状态TTL:");
            System.out.println("   - 避免状态无限增长");
            System.out.println("   - 根据业务需求设置合理的过期时间");
            System.out.println("   - 考虑状态访问模式");
            
            System.out.println("\n3. 🔧 优化RocksDB配置:");
            System.out.println("   - 调整块缓存大小");
            System.out.println("   - 配置写缓冲区");
            System.out.println("   - 设置合适的压缩算法");
            
            System.out.println("\n4. 📸 Checkpoint优化:");
            System.out.println("   - 设置合理的间隔时间");
            System.out.println("   - 避免过多的并发Checkpoint");
            System.out.println("   - 使用增量Checkpoint");
            
            System.out.println("\n5. 🧹 状态清理:");
            System.out.println("   - 定期清理过期状态");
            System.out.println("   - 监控状态大小");
            System.out.println("   - 实现状态压缩");
            
            System.out.println("========================");
        }
    }
    
    /**
     * 运行状态管理工具
     */
    public static void runStateManagementTools(StreamExecutionEnvironment env) {
        System.out.println("\n=== 状态管理工具 ===");
        
        // 验证状态后端配置
        StateManagementUtils.validateStateBackendConfig(env);
        
        // 打印最佳实践建议
        StateManagementUtils.printStateManagementTips();
    }
} 